import math
import numpy as np

def anint(x, lx):
    lhx=0.5*lx;
    while x>lhx:
        x-=lx
    while x<-lhx:
        x+=lx
    return x

def distancia(lbox,pos1,pos2):
    dis2=0
    drx=pos2[0]-pos1[0]
    drx=anint(drx,lbox[0])
    dry=pos2[1]-pos1[1]
    dry=anint(dry,lbox[1])
    drz=pos2[2]-pos1[2]
    drz=anint(drz,lbox[2])
    dis2=drx*drx+dry*dry+drz*drz
    return  [math.sqrt(dis2),drx,dry,drz]

def BCD(lbox,pos1):
    while pos1[0]>lbox[0]:
        pos1[0]-=lbox[0]
    while pos1[0]<0:
        pos1[0]+=lbox[0]
    while pos1[1]>lbox[1]:
        pos1[1]-=lbox[1]
    while pos1[1]<0:
        pos1[1]+=lbox[1]
    while pos1[2]>lbox[2]:
        pos1[2]-=lbox[2]
    while pos1[2]<0:
        pos1[2]+=lbox[2]
    return pos1
    
def importaDatos(name):
    pos=[]
    diametro=[]
    count=0
    
    for line in file(name):
        line=line.replace('\n','')
        val=line.split(' ')
        if count!=0:
            res=[]
            for v in val:
                if v!='':
                    res.append(float(v))
            pos.append([res[0],res[1],res[2]])
            diametro.append(res[9])
        else:
            lx=float(val[8])
            ly=float(val[9])
            lz=float(val[10])
            lbox=[lx,ly,lz]
        count+=1
    
    return [lbox,pos,diametro]

def calcDistancias(pos,lbox):
    distancias=[]
    for i in range(len(pos)):
        d1=[]
        for j in range(i,len(pos)):
            dis = distancia(lbox,pos[i],pos[j])
            d1.append(dis)
        distancias.append(d1)
    return distancias

def calcZeta(dist,beta):
    zetaxy=[]
    zetaxz=[]
    zetayz=[]
    for i in range(len(dist)):
        auxxy=[]
        auxxz=[]
        auxyz=[]
        for j in range(i,len(dist)):
            r5=dist[i][j-i][0]
            if r5>1.e-4:
                r5=r5*r5*r5*r5*r5
                auxxy.append(3.*dist[i][j-i][1]*dist[i][j-i][2]/r5*beta)
                auxxz.append(3.*dist[i][j-i][1]*dist[i][j-i][3]/r5*beta)
                auxyz.append(3.*dist[i][j-i][2]*dist[i][j-i][3]/r5*beta)
            else:
                auxxy.append(0.0)
                auxxz.append(0.0)
                auxyz.append(0.0)
        zetaxy.append(auxxy)
        zetaxz.append(auxxz)
        zetayz.append(auxyz)
    return [zetaxy,zetaxz,zetayz]

def calcChi(dist,beta):
    chix=[]
    chiy=[]
    chiz=[]
    for i in range(len(dist)):
        #Componente ii
        auxx=[]
        auxy=[]
        auxz=[]
        for j in range(i,len(dist)):
            r=dist[i][j-i][0]
            if r>1.e-8:
                r2=r*r
                r5=r2*r2*r
                auxx.append((3.*dist[i][j-i][1]*dist[i][j-i][1]-r2)/r5*beta)
                auxy.append((3.*dist[i][j-i][2]*dist[i][j-i][2]-r2)/r5*beta)
                auxz.append((3.*dist[i][j-i][3]*dist[i][j-i][3]-r2)/r5*beta)
            else:
                auxx.append(-1.0)
                auxy.append(-1.0)
                auxz.append(-1.0)
        chix.append(auxx)
        chiy.append(auxy)
        chiz.append(auxz)
    return [chix,chiy,chiz]

def construyeMatriz(zetaxy,zetaxz,zetayz,chix,chiy,chiz):
    npart=len(zetaxy[0])
    matriz=np.zeros(((3*npart,3*npart)))  

    for i in range(npart):
        aux=[]
        for j in range(npart):
            idi=i
            idj=j
            if j < i:
                idi=j
                idj=i
            matriz[i][j]=-chix[idi][idj-idi]
        for j in range(npart):
            idi=i
            idj=j
            if j < i:
                idi=j
                idj=i
            matriz[i][npart+j]=-zetaxy[idi][idj-idi]
        for j in range(npart):
            idi=i
            idj=j
            if j < i:
                idi=j
                idj=i
            matriz[i][2*npart+j]=-zetaxz[idi][idj-idi]     
        
    for i in range(npart):
        aux=[]
        for j in range(npart):
            idi=i
            idj=j
            if j < i:
                idi=j
                idj=i
            matriz[i+npart][j]=-zetaxy[idi][idj-idi]
        for j in range(npart):
            idi=i
            idj=j
            if j < i:
                idi=j
                idj=i
            matriz[i+npart][j+npart]=-chiy[idi][idj-idi]
        for j in range(npart):
            idi=i
            idj=j
            if j < i:
                idi=j
                idj=i
            matriz[i+npart][j+2*npart]=-zetaxz[idi][idj-idi]
        
    for i in range(npart):
        for j in range(npart):
            idi=i
            idj=j
            if j < i:
                idi=j
                idj=i
            matriz[i+2*npart][j]=-zetaxz[idi][idj-idi]
        for j in range(npart):
            idi=i
            idj=j
            if j < i:
                idi=j
                idj=i
            matriz[i+2*npart][j+npart]=-zetaxy[idi][idj-idi]
        for j in range(npart):
            idi=i
            idj=j
            if j < i:
                idi=j
                idj=i
            matriz[i+2*npart][j+2*npart]=-chiz[idi][idj-idi]

    
    for i in range(3*npart):
        for j in range(i+1,3*npart):
            matriz[j][i]=matriz[i][j]
    return matriz

                    

name='system_bk.dat'
beta=1.0/8.0

[lbox,pos,diametro] = importaDatos(name)
npart=len(pos)
distancias = calcDistancias(pos,lbox)
[zetaxy,zetaxz,zetayz] = calcZeta(distancias,beta)

[chix,chiy,chiz] = calcChi(distancias,beta)

matriz = construyeMatriz(zetaxy,zetaxz,zetayz,chix,chiy,chiz)
cont = np.zeros((3*npart,1))
for i in range(2*npart,3*npart):
    cont[i][0]=1.0

solucion = np.linalg.solve (matriz , cont )
f = open('salida.dat','w')
for i in range(npart):
    f.write(str(solucion[i])+' ')
    f.write(str(solucion[i+npart])+' ')
    f.write(str(solucion[i+2*npart])+'\n')
f.close()  

f = open('zeta.dat','w')
for i in range(npart):
    for j in range(i,npart):
        f.write(str(zetaxy[i][j-i])+' ')
    f.write('\n')
f.close()   

f = open('chi.dat','w')
for i in range(npart):
    for j in range(i,npart):
        f.write(str(chix[i][j-i])+' ')
    f.write('\n')
f.close()   

f = open('matriz.dat','w')
for i in range(3*npart):
    for j in range(3*npart):
        f.write(str(matriz[i][j])+' ')
    f.write('\n')
f.close()   

iii=0
jjj=1
print 'part(',iii,'): ',pos[iii]
print 'part(',jjj,'): ',pos[jjj]
[r,dx,dy,dz]=distancia(lbox,pos[iii],pos[jjj])
print "Test funcion zeta"
print 3.*beta*dx*dy/math.pow(r, 5),3.*beta*dx*dz/math.pow(r, 5),3.*beta*dz*dy/math.pow(r, 5)
print zetaxy[iii][jjj-iii],zetaxz[iii][jjj-iii],zetayz[iii][jjj-iii]
print "Test funcion chi"
print beta*(3.*dx*dx-r*r)/math.pow(r, 5),beta*(3.*dy*dy-r*r)/math.pow(r, 5),beta*(3.*dz*dz-r*r)/math.pow(r, 5)
print chix[iii][jjj-iii],chiy[iii][jjj-iii],chiz[iii][jjj-iii]

for i in range(npart):
    print solucion[i],solucion[i+npart],solucion[i+2*npart]
